# -*- coding: utf-8 -*-
from __future__ import division, print_function

import glob
import os
import time

import cv2
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import selectivesearch
import skimage.data
from tqdm import tqdm
import pickle

def main():

    image_set_trainval_path = "../_DATASET/NWPU_VHR/ImageSets/trainval.txt"
    image_set_test_path = "../_DATASET/NWPU_VHR/ImageSets/test.txt"
    
    pickle_trainval_path = 'data/selective_search_data/nwpu_trainval.pkl'
    pickle_test_path = 'data/selective_search_data/nwpu_test.pkl'
    
    #phase = ['trainval','test']
    
    phase = ['test']
    save_vis = False
    save_dir = "images/NWPU/"
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
        
    imgs_dir = "../_DATASET/NWPU_VHR/AllImages/"

    img_path_list = glob.glob(imgs_dir + "*.jpg")
    candidate_list = []
    score_list = []
    for p in phase:
        print('generate {} proposals'.format(p))
        if p == "trainval":
            pickle_save_path = pickle_trainval_path
            image_set_path = image_set_trainval_path
            index_list_train = []
            index_list = index_list_train
        elif p == "test":
            pickle_save_path = pickle_test_path
            image_set_path = image_set_test_path
            index_list_test = []
            index_list = index_list_test
        f = open(image_set_path,'r')
        image_set_list = [str(x.strip()) for x in f.readlines()]
        print(image_set_list)
        
        for img_path in tqdm(img_path_list):
            
            index = os.path.basename(img_path).split(".")[0]
    #         print(index)
            if str(index) in image_set_list:
                index_list.append(str(index))
        #         print(index_list)
                img = cv2.imread(img_path)

                # im_orig:输入图片
                # scale：表示felzenszwalb分割时，值越大，表示保留的下来的集合就越大
                # sigma：表示felzenszwalb分割时，用的高斯核宽度
                # min_size：表示分割后最小组尺寸
        #         start = time.time()
                img_lbl, regions = selectivesearch.selective_search(img,scale=200,sigma=0.8,min_size=50)
        #         end = time.time()
        #         print(end-start)

                candidate_per_img_list = []
                score_per_img_list = []
                for region in tqdm(regions):
                    candidate_per_img_list.append(region['rect'])
                    score_per_img_list.append(1)

                if save_vis:
                    fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(10, 10))
                    ax.imshow(img[:, :, ::-1])
                    for x, y, w, h in tqdm(candidate_per_img_list):
                        #         print(x, y, w, h)
                        rect = mpatches.Rectangle((x, y), w, h, fill=False, edgecolor="red", linewidth=1)
                        ax.add_patch(rect)
                    plt.show()
                    fig.savefig(os.path.join(save_dir,index + '.jpg'), dpi=50)    
                    
                candidate_per_img_set = set(candidate_per_img_list)
                candidate_list.append(np.array(candidate_per_img_set))
                score_list.append(np.array(score_per_img_list)) 
            else:
                continue

        pickle_dict = {'indexes':index_list,
                       'boxes':candidate_list,
                       'scores':score_list}

    #     print(len(candidate_list))
    #     print(len(score_list))
        f=open(pickle_save_path,'wb')
        pickle.dump(pickle_dict,f)
        f.close()
    

if __name__ == "__main__":
    main()